// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.

// ************************************************************************* //
//  File: avtPDFFilter.C
// ************************************************************************* //

#include <avtPDFFilter.h>

#include <vtkCellData.h>
#include <vtkCellDataToPointData.h>
#include <vtkFloatArray.h>
#include <vtkPointData.h>
#include <vtkPointDataToCellData.h>
#include <vtkRectilinearGrid.h>

#include <avtExtents.h>
#include <avtParallel.h>
#include <avtTypes.h>

#include <ImproperUseException.h>
#include <InvalidVariableException.h>


// ****************************************************************************
//  Method: avtPDFFilter constructor
//
//  Programmer: childs -- generated by xml2avt
//  Creation:   Sun Nov 20 12:22:59 PDT 2005
//
// ****************************************************************************

avtPDFFilter::avtPDFFilter()
{
}


// ****************************************************************************
//  Method: avtPDFFilter destructor
//
//  Programmer: childs -- generated by xml2avt
//  Creation:   Sun Nov 20 12:22:59 PDT 2005
//
//  Modifications:
//
// ****************************************************************************

avtPDFFilter::~avtPDFFilter()
{
}


// ****************************************************************************
//  Method:  avtPDFFilter::Create
//
//  Programmer: childs -- generated by xml2avt
//  Creation:   Sun Nov 20 12:22:59 PDT 2005
//
// ****************************************************************************

avtFilter *
avtPDFFilter::Create()
{
    return new avtPDFFilter();
}


// ****************************************************************************
//  Method:      avtPDFFilter::SetAtts
//
//  Purpose:
//      Sets the state of the filter based on the attribute object.
//
//  Arguments:
//      a        The attributes to use.
//
//  Programmer: childs -- generated by xml2avt
//  Creation:   Sun Nov 20 12:22:59 PDT 2005
//
// ****************************************************************************

void
avtPDFFilter::SetAtts(const AttributeGroup *a)
{
    atts = *(const PDFAttributes*)a;
}


// ****************************************************************************
//  Method: avtPDFFilter::Equivalent
//
//  Purpose:
//      Returns true if creating a new avtPDFFilter with the given
//      parameters would result in an equivalent avtPDFFilter.
//
//  Programmer: childs -- generated by xml2avt
//  Creation:   Sun Nov 20 12:22:59 PDT 2005
//
// ****************************************************************************

bool
avtPDFFilter::Equivalent(const AttributeGroup *a)
{
    return (atts == *(PDFAttributes*)a);
}


// ****************************************************************************
//  Method: avtPDFFilter::PreExecute
//
//  Purpose:
//      Called before "Execute".  This will set up the range's we are operating
//      on and the data structures we use to count the density.
//
//  Programmer: Hank Childs
//  Creation:   November 21, 2005
//
//  Modifications:
//    Jeremy Meredith, Thu Feb 15 11:55:03 EST 2007
//    Call inherited PreExecute before everything else.
//
// ****************************************************************************

void
avtPDFFilter::PreExecute(void)
{
    avtPluginDataTreeIterator::PreExecute();

    //
    // Make sure that all the variables we need are in the input.
    //
    avtDataAttributes &d_atts = GetInput()->GetInfo().GetAttributes();
    if (! d_atts.ValidVariable(atts.GetVar1()))
        EXCEPTION1(InvalidVariableException, atts.GetVar1());
    if (! d_atts.ValidVariable(atts.GetVar2()))
        EXCEPTION1(InvalidVariableException, atts.GetVar2());
    if (atts.GetNumAxes() == PDFAttributes::Three)
        if (! d_atts.ValidVariable(atts.GetVar3()))
            EXCEPTION1(InvalidVariableException, atts.GetVar3());

    //
    // Now calculate what the extents of the variables are.
    //
    bool needVar1Exts = !(atts.GetVar1MinFlag()) || !(atts.GetVar1MaxFlag());
    if (needVar1Exts)
    {
        double exts[2];
        GetDataExtents(exts, atts.GetVar1().c_str());
        var1Min = exts[0];
        var1Max = exts[1];
    }
    if (atts.GetVar1MinFlag())
        var1Min = atts.GetVar1Min();
    if (atts.GetVar1MaxFlag())
        var1Max = atts.GetVar1Max();
    bool needVar2Exts = !(atts.GetVar2MinFlag()) || !(atts.GetVar2MaxFlag());
    if (needVar2Exts)
    {
        double exts[2];
        GetDataExtents(exts, atts.GetVar2().c_str());
        var2Min = exts[0];
        var2Max = exts[1];
    }
    if (atts.GetVar2MinFlag())
        var2Min = atts.GetVar2Min();
    if (atts.GetVar2MaxFlag())
        var2Max = atts.GetVar2Max();
    if (atts.GetNumAxes() == PDFAttributes::Three)
    {
        bool needVar3Exts = !(atts.GetVar3MinFlag())||!(atts.GetVar3MaxFlag());
        if (needVar3Exts)
        {
            double exts[2];
            GetDataExtents(exts, atts.GetVar3().c_str());
            var3Min = exts[0];
            var3Max = exts[1];
        }
        if (atts.GetVar3MinFlag())
            var3Min = atts.GetVar3Min();
        if (atts.GetVar3MaxFlag())
            var3Max = atts.GetVar3Max();
    }

    //
    // Now set up the array that will let us track the density.
    //
    if (atts.GetVar1NumSamples() <= 1)
        EXCEPTION0(ImproperUseException);
    if (atts.GetVar2NumSamples() <= 1)
        EXCEPTION0(ImproperUseException);
    if (atts.GetNumAxes() == PDFAttributes::Three &&
        atts.GetVar3NumSamples() <= 1)
        EXCEPTION0(ImproperUseException);

    int numSamples = atts.GetVar1NumSamples() * atts.GetVar2NumSamples();
    if (atts.GetNumAxes() == PDFAttributes::Three)
        numSamples *= atts.GetVar3NumSamples();
    density = new int[numSamples];
    for (int i = 0 ; i < numSamples ; i++)
        density[i] = 0;
}


// ****************************************************************************
//  Method: avtPDFFilter::PostExecute
//
//  Purpose:
//      Called after "Execute".  This will set up the final probability 
//      density function on a rectilinear grid.
//
//  Programmer: Hank Childs
//  Creation:   November 21, 2005
//
// ****************************************************************************

void
avtPDFFilter::PostExecute(void)
{
    int   i;

    //
    // Start off by setting up a rectilinear grid.
    //
    vtkRectilinearGrid *rgrid = vtkRectilinearGrid::New();
    
    int nX = atts.GetVar1NumSamples()+1;
    vtkFloatArray *x = vtkFloatArray::New();
    x->SetNumberOfTuples(nX);
    double minX, maxX;
    if (atts.GetScaleCube())
    {
        minX = 0.;
        maxX = 1.;
    }
    else
    {
        minX = var1Min;
        maxX = var1Max;
    }
    double stepX = (maxX-minX) / (nX - 1);
    for (i = 0 ; i < nX ; i++)
        x->SetTuple1(i, minX + i*stepX);
    rgrid->SetXCoordinates(x);
    x->Delete();

    int nY = atts.GetVar2NumSamples()+1;
    vtkFloatArray *y = vtkFloatArray::New();
    y->SetNumberOfTuples(nY);
    double minY, maxY;
    if (atts.GetScaleCube())
    {
        minY = 0.;
        maxY = 1.;
    }
    else
    {
        minY = var2Min;
        maxY = var2Max;
    }
    double stepY = (maxY-minY) / (nY - 1);
    for (i = 0 ; i < nY ; i++)
        y->SetTuple1(i, minY + i*stepY);
    rgrid->SetYCoordinates(y);
    y->Delete();

    vtkFloatArray *z = vtkFloatArray::New();
    int nZ;
    if (atts.GetNumAxes() == PDFAttributes::Three)
    {
        nZ = atts.GetVar3NumSamples()+1;
        z->SetNumberOfTuples(nZ);
        double minZ, maxZ;
        if (atts.GetScaleCube())
        {
            minZ = 0.;
            maxZ = 1.;
        }
        else
        {
            minZ = var3Min;
            maxZ = var3Max;
        }
        double stepZ = (maxZ-minZ) / (nZ - 1);
        for (i = 0 ; i < nZ ; i++)
            z->SetTuple1(i, minZ + i*stepZ);
    }
    else
    {
        nZ = 1;
        z->SetNumberOfTuples(1);
        z->SetTuple1(0, 0.);
    }
    rgrid->SetZCoordinates(z);
    z->Delete();

    rgrid->SetDimensions(nX, nY, nZ);

    //
    // We have the density for this processor.  Get it for all processors.
    //
    int numSamples = atts.GetVar1NumSamples() * atts.GetVar2NumSamples();
    if (atts.GetNumAxes() == PDFAttributes::Three)
        numSamples *= atts.GetVar3NumSamples();
    int *density2 = new int[numSamples];
    SumIntArrayAcrossAllProcessors(density, density2, numSamples);

    //
    // Now set up the output variable.
    //
    vtkFloatArray *densityV = vtkFloatArray::New();
    densityV->SetName(pipelineVariable);
    densityV->SetNumberOfTuples(numSamples);
    double denom = 1.;
    if (atts.GetDensityType() == PDFAttributes::Probability)
    {
        denom = 0.;
        for (i = 0 ; i < numSamples ; i++)
            denom += density2[i];
        if (denom == 0.)
            denom = 1.;
    }
    for (i = 0 ; i < numSamples ; i++)
        densityV->SetTuple1(i, density2[i] / denom);
    rgrid->GetCellData()->SetScalars(densityV);
    densityV->Delete();

    //
    // Now put it back in the world of AVT.
    //
    avtDataTree_p tree = new avtDataTree();
    if (PAR_Rank() == 0)
        tree = new avtDataTree(rgrid, 0);
    rgrid->Delete();
    SetOutputDataTree(tree);

    delete [] density;
    density = NULL;
    delete [] density2;

    //
    // Call this one more time.  Now we have all of the info for it to set up
    // the extents properly.
    //
    UpdateDataObjectInfo();
}


// ****************************************************************************
//  Method: avtPDFFilter::UpdateDataObjectInfo
//
//  Purpose:
//      Modifies the data object information to reflect the new, output mesh.
//
//  Programmer: Hank Childs
//  Creation:   November 21, 2005
//
//  Modifications:
//
//    Hank Childs, Thu Aug 26 13:47:30 PDT 2010
//    Change extents names.
//
//    Brad Whitlock, Mon Apr  7 15:55:02 PDT 2014
//    Add filter metadata used in export.
//    Work partially supported by DOE Grant SC0007548.
//
// ****************************************************************************

void
avtPDFFilter::UpdateDataObjectInfo(void)
{
    avtDataAttributes &in_atts  = GetInput()->GetInfo().GetAttributes();
    avtDataAttributes &out_atts = GetOutput()->GetInfo().GetAttributes();
    if (pipelineVariable != NULL)
    {
        int nV = in_atts.GetNumberOfVariables();
        for (int i = 0 ; i < nV ; i++)
        {
            const std::string &name = in_atts.GetVariableName(i);
            out_atts.RemoveVariable(name);
        }
        out_atts.AddVariable(pipelineVariable);
        out_atts.SetVariableDimension(1, pipelineVariable);
        out_atts.SetVariableType(AVT_SCALAR_VAR, pipelineVariable);
        out_atts.SetCentering(AVT_ZONECENT, pipelineVariable);
        out_atts.SetActiveVariable(pipelineVariable);
    }

    int dim = (atts.GetNumAxes() == PDFAttributes::Three ? 3 : 2);
    out_atts.SetTopologicalDimension(dim);
    out_atts.SetSpatialDimension(dim);

    if (pipelineVariable != NULL)
    {
        double exts[6] = { 0, 1, 0, 1, 0, 1 };
        if (!atts.GetScaleCube())
        {
            exts[0] = var1Min;
            exts[1] = var1Max;
            exts[2] = var2Min;
            exts[3] = var2Max;
            if (dim == 3)
            {
                exts[4] = var3Min;
                exts[5] = var3Max;
            }
        }
        if (dim == 2)
            exts[4] = exts[5] = 0.;
        out_atts.GetOriginalSpatialExtents()->Set(exts);
        out_atts.GetThisProcsOriginalSpatialExtents()->Set(exts);
        out_atts.GetDesiredSpatialExtents()->Set(exts);
        out_atts.GetActualSpatialExtents()->Set(exts);
        out_atts.GetThisProcsActualSpatialExtents()->Set(exts);

        out_atts.SetXLabel(atts.GetVar1());
        out_atts.SetXUnits(in_atts.GetVariableUnits(atts.GetVar1().c_str()));
        out_atts.SetYLabel(atts.GetVar2());
        out_atts.SetYUnits(in_atts.GetVariableUnits(atts.GetVar2().c_str()));
        if (dim == 3)
        {
            out_atts.SetZLabel(atts.GetVar3());
            out_atts.SetZUnits(in_atts.GetVariableUnits(atts.GetVar3().c_str()));
        }
    }
    out_atts.AddFilterMetaData("PDF");
}


// ****************************************************************************
//  Method: avtPDFFilter::ModifyContract
//
//  Purpose:
//      States the secondary variables we need to make the PDF.
//
//  Programmer: Hank Childs
//  Creation:   November 21, 2005
//
// ****************************************************************************

avtContract_p
avtPDFFilter::ModifyContract(avtContract_p s)
{
    avtContract_p spec = new avtContract(s);
    avtDataRequest_p dataRequest = spec->GetDataRequest();
    dataRequest->AddSecondaryVariable(atts.GetVar1().c_str());
    dataRequest->AddSecondaryVariable(atts.GetVar2().c_str());
    if (atts.GetNumAxes() == PDFAttributes::Three)
        dataRequest->AddSecondaryVariable(atts.GetVar3().c_str());

    return spec;
}


// ****************************************************************************
//  Method: avtPDFFilter::ExecuteData
//
//  Purpose:
//      Sends the specified input and output through the PDF filter.
//
//  Arguments:
//      in_dr      The input data representation.
//
//  Returns:       The output data representation.
//
//  Programmer: childs -- generated by xml2avt
//  Creation:   Sun Nov 20 12:22:59 PDT 2005
//
//  Modifications:
//    Eric Brugger, Thu Jul 31 12:09:01 PDT 2014
//    Modified the class to work with avtDataRepresentation.
//
// ****************************************************************************

avtDataRepresentation *
avtPDFFilter::ExecuteData(avtDataRepresentation *in_dr)
{
    //
    // Get the VTK data set.
    //
    vtkDataSet *in_ds = in_dr->GetDataVTK();

    bool isNodal1 = true;
    vtkDataArray *arr1=in_ds->GetPointData()->GetArray(atts.GetVar1().c_str());
    if (arr1 == NULL)
    {
        arr1 = in_ds->GetCellData()->GetArray(atts.GetVar1().c_str());
        isNodal1 = false;
    }
    if (arr1 == NULL)
    {
        EXCEPTION0(ImproperUseException);
    }

    bool isNodal2 = true;
    vtkDataArray *arr2=in_ds->GetPointData()->GetArray(atts.GetVar2().c_str());
    if (arr2 == NULL)
    {
        arr2 = in_ds->GetCellData()->GetArray(atts.GetVar2().c_str());
        isNodal2 = false;
    }
    if (arr2 == NULL)
    {
        EXCEPTION0(ImproperUseException);
    }
    bool deleteArr2 = false;
    if (isNodal2 != isNodal1)
    {
        if (isNodal2) // 2 is nodal, 1 is zonal.  So convert 2 to zonal
        {
            vtkPointDataToCellData *pd2cd = vtkPointDataToCellData::New();
            vtkDataSet *new_in = (vtkDataSet *) in_ds->NewInstance();
            new_in->CopyStructure(in_ds);
            new_in->GetPointData()->AddArray(
                      in_ds->GetPointData()->GetArray(atts.GetVar2().c_str()));
            pd2cd->SetInputData(new_in);
            pd2cd->Update();
            arr2 = pd2cd->GetOutput()->GetCellData()
                                            ->GetArray(atts.GetVar2().c_str());
            arr2->Register(NULL);
            pd2cd->Delete();
            new_in->Delete();
            deleteArr2 = true;
        }
        else // 2 is zonal, 1 is nodal.  So convert 2 to nodal
        {
            vtkCellDataToPointData *cd2pd = vtkCellDataToPointData::New();
            vtkDataSet *new_in = (vtkDataSet *) in_ds->NewInstance();
            new_in->CopyStructure(in_ds);
            new_in->GetCellData()->AddArray(
                       in_ds->GetCellData()->GetArray(atts.GetVar2().c_str()));
            cd2pd->SetInputData(new_in);
            cd2pd->Update();
            arr2 = cd2pd->GetOutput()->GetPointData()
                                           ->GetArray(atts.GetVar2().c_str());
            arr2->Register(NULL);
            cd2pd->Delete();
            new_in->Delete();
            deleteArr2 = true;
        }
    }

    bool is3D = (atts.GetNumAxes() == PDFAttributes::Three);
    bool isNodal3 = isNodal2;
    vtkDataArray *arr3 = NULL;
    bool deleteArr3 = false;
    if (is3D)
    {
        isNodal3 = true;
        arr3 = in_ds->GetPointData()->GetArray(atts.GetVar3().c_str());
        if (arr3 == NULL)
        {
            arr3 = in_ds->GetCellData()->GetArray(atts.GetVar3().c_str());
            isNodal3 = false;
        }
        if (arr3 == NULL)
        {
            EXCEPTION0(ImproperUseException);
        }
        if (isNodal3 != isNodal1)
        {
            if (isNodal3) // 3 is nodal, 1 is zonal.  So convert 3 to zonal
            {
                vtkPointDataToCellData *pd2cd = vtkPointDataToCellData::New();
                vtkDataSet *new_in = (vtkDataSet *) in_ds->NewInstance();
                new_in->CopyStructure(in_ds);
                new_in->GetPointData()->AddArray(
                      in_ds->GetPointData()->GetArray(atts.GetVar3().c_str()));
                pd2cd->SetInputData(new_in);
                pd2cd->Update();
                arr3 = pd2cd->GetOutput()->
                               GetCellData()->GetArray(atts.GetVar3().c_str());
                arr3->Register(NULL);
                pd2cd->Delete();
                new_in->Delete();
                deleteArr3 = true;
            }
            else // 3 is zonal, 1 is nodal.  So convert 3 to nodal
            {
                vtkCellDataToPointData *cd2pd = vtkCellDataToPointData::New();
                vtkDataSet *new_in = (vtkDataSet *) in_ds->NewInstance();
                new_in->CopyStructure(in_ds);
                new_in->GetCellData()->AddArray(
                       in_ds->GetCellData()->GetArray(atts.GetVar3().c_str()));
                cd2pd->SetInputData(new_in);
                cd2pd->Update();
                arr3 = cd2pd->GetOutput()->
                              GetPointData()->GetArray(atts.GetVar3().c_str());
                arr3->Register(NULL);
                cd2pd->Delete();
                new_in->Delete();
                deleteArr3 = true;
            }
        }
    }

    int nvals = arr1->GetNumberOfTuples();
    int nX = atts.GetVar1NumSamples();
    float xStep = (var1Max-var1Min) / nX;
    int nY = atts.GetVar2NumSamples();
    float yStep = (var2Max-var2Min) / nY;
    int nZ = atts.GetVar3NumSamples();
    if (nZ <= 0)
        nZ = 1;
    float zStep = (var3Max-var3Min) / nZ;
    for (int i = 0 ; i < nvals ; i++)
    {
        float val1 = arr1->GetTuple1(i);
        float val2 = arr2->GetTuple1(i);
        float val3 = 0.;
        if (is3D)
            val3 = arr3->GetTuple1(i);
        int index = 0;
        int idxX = (int)((val1-var1Min) / xStep);
        if (idxX >= nX)
            idxX = nX-1;
        if (idxX < 0)
            idxX = 0;
        index += idxX;
        int idxY = (int)((val2-var2Min) / yStep);
        if (idxY >= nY)
            idxY = nY-1;
        if (idxY < 0)
            idxY = 0;
        idxY *= nX;
        index += idxY;
        if (is3D)
        {
            int idxZ = (int)((val3-var3Min) / zStep);
            if (idxZ >= nZ)
                idxZ = nZ-1;
            if (idxZ < 0)
                idxZ = 0;
            idxZ *= (nX*nY);
            index += idxZ;
        }
        density[index] += 1;
    }

    if (deleteArr2)
        arr2->Delete();
    if (deleteArr3)
        arr3->Delete();

    //
    // We didn't really construct a new data set.  So just return NULL.
    //
    return NULL;
}


